
import asyncio
import subprocess

import numpy
import ffmpeg


class LJFfmpeg():
    def __init__(self, rtsp_url, websocket=None) -> None:
        print(websocket)
        self.websocket = websocket
        self.rtsp_url = rtsp_url
        self.is_stop = False

    async def start(self):

        print(f"LJ - 开始执行视频流[{self.rtsp_url}]转发...")
        probe = ffmpeg.probe(self.rtsp_url)
        video_stream = next(
            (stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)
        width = int(video_stream['width'])
        height = int(video_stream['height'])

        # ffmpeg_process = (
        #     ffmpeg
        #     .input(self.rtsp_url, rtsp_transport='tcp')
        #     .output('pipe:', format='rawvideo', pix_fmt='bgr24', loglevel="quiet", r=25)
        #     .run_async(pipe_stdout=True)
        # )

        ffmpeg_process = subprocess.Popen([
            'ffmpeg',
            '-rtsp_transport', 'tcp',  # 使用 TCP 连接 RTSP
            '-i', self.rtsp_url,  # 替换为你的 RTSP 视频流的 URL
            '-f', 'mpegts',  # 输出为 MPEG-TS 流
            '-codec:v', 'mpeg1video',  # 使用 MPEG-1 视频编解码器
            # '-s', '640x480',  # 视频分辨率
            '-b:v', '1000k',  # 视频比特率
            '-bf', '0',  # 禁用帧间预测
            '-r', '30',  # 视频帧率
            '-',
        ], stdout=subprocess.PIPE, bufsize=10**8)

        cnt_empty = 0

        while not self.is_stop:
            in_bytes = ffmpeg_process.stdout.read(height * width * 3)
            if not in_bytes:
                cnt_empty += 1
                if cnt_empty > 10:
                    break
                continue
            cnt_empty = 0

            await self.websocket.send(in_bytes)
            # import cv2
            # image = numpy.frombuffer(
            #     in_bytes, dtype='uint8').reshape((height, width, 3))
            # cv2.imshow('RTSP Stream', image)
            # if cv2.waitKey(1) & 0xFF == ord('q'):
            #     break

        # 在完成后释放资源
        # cv2.destroyAllWindows()
        ffmpeg_process.kill()

    def stop(self):
        self.is_stop = True
